import time
import threading
import numpy as np
import pandas as pd
import sqlite3
from tqsdk import TqApi, TqAuth, TqKq
from MyTT import EMA, SUM
from rich.console import Console
from rich.theme import Theme
from MyTT import REF, EMA

# 定义颜色主题
custom_theme = Theme({
    "info": "cyan",
    "warning": "yellow",
    "error": "red",
    "success": "green",
    "trade": "bold magenta"
})
console = Console(theme=custom_theme)


class TradingBot:
    def __init__(self):
        self.api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba"))
        self.running = True
        self.threads = []
        self.lock = threading.Lock()
        self.console = console

    def get_monitored_symbols(self):
        """从数据库获取监控品种"""
        try:
            conn = sqlite3.connect('trade_monitor.db')
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM products WHERE is_monitored=1")
            products = cursor.fetchall()
            conn.close()
            return [row[2] for row in products]
        except Exception as e:
            self.console.print(f"从数据库获取监控品种时出错: {e}", style="error")
            return []

    def start_trading(self):
        """启动交易"""
        symbols = self.get_monitored_symbols()
        if not symbols:
            self.console.print("未获取到监控品种，程序退出", style="warning")
            return

        self.console.print(f"开始监控品种: {symbols}", style="info")

        # 为每个品种创建交易线程
        for symbol in symbols:
            thread = threading.Thread(
                target=self.trade_loop,
                args=(symbol,),
                daemon=True
            )
            thread.start()
            self.threads.append(thread)

        # 主线程保持运行
        try:
            while self.running:
                time.sleep(1)
        except KeyboardInterrupt:
            self.console.print("\n接收到中断信号，程序准备停止", style="warning")
        finally:
            self.stop_trading()

    def stop_trading(self):
        """停止交易"""
        self.running = False
        for thread in self.threads:
            thread.join()
        if self.api:
            self.api.close()

    def trade_loop(self, symbol):
        """单个品种的交易循环"""
        try:
            quote = self.api.get_quote(symbol)
            klines = self.api.get_kline_serial(symbol, 60, data_length=300)

            # 注册更新通知
            update_event = threading.Event()
            self.api.register_update_notify(update_event.set)

            while self.running:
                # 等待更新事件或超时
                update_event.wait(timeout=1)
                update_event.clear()

                if self.api.is_changing(quote):
                    self.check_stoploss(symbol, quote, klines)

        except Exception as e:
            self.console.print(f"处理品种 {symbol} 时出错: {e}", style="error")

    def check_stoploss(self, symbol, quote, klines):
        """止损逻辑"""
        try:
            close = klines.close
            low = klines.low
            open = klines.open
            high = klines.high

            if close is None or low is None or open_price is None or high is None:
                return

            # 转换为列表，方便使用 MyTT 计算
            close = list(close)
            low = list(low)
            open_price = list(open_price)
            high = list(high)

            Q=(3*close+low+open+high)/6

            # 计算操盘线
            trading_line=(26*Q+25*REF(Q,1)+24*REF(Q,2)+23*REF(Q,3)+22*REF(Q,4)+21*REF(Q,5)+20*REF(Q,6)+19*REF(Q,7)+18*REF(Q,8)+17*REF(Q,9)+16*REF(Q,10)+15*REF(Q,11)+14*REF(Q,12)+13*REF(Q,13)+12*REF(Q,14)+11*REF(Q,15)+10*REF(Q,16)+9*REF(Q,17)+8*REF(Q,18)+7*REF(Q,19)+6*REF(Q,20)+5*REF(Q,21)+4*REF(Q,22)+3*REF(Q,23)+2*REF(Q,24)+REF(Q,26))/351

            # 计算空头线
            short_line = EMA(trading_line, 7)

            # 使用锁保证线程安全
            with self.lock:
                position = self.api.get_position(symbol)

                # 平仓条件判断
                if len(trading_line) >= 2 and len(short_line) >= 2:
                    if position.pos_short > 0 and trading_line[-2] > short_line[-2]:
                        self.api.insert_order(
                            symbol=symbol,
                            direction="BUY",
                            offset="CLOSE",
                            volume=position.pos_short
                        )
                        self.console.print(
                            f"{time.strftime('%X')} 平空仓: {symbol} 手数: {position.pos_short}",
                            style="trade"
                        )

                    if position.pos_long > 0 and trading_line[-2] < short_line[-2]:
                        self.api.insert_order(
                            symbol=symbol,
                            direction="SELL",
                            offset="CLOSE",
                            volume=position.pos_long
                        )
                        self.console.print(
                            f"{time.strftime('%X')} 平多仓: {symbol} 手数: {position.pos_long}",
                            style="trade"
                        )

        except Exception as e:
            self.console.print(f"处理 {symbol} 止损逻辑时出错: {e}", style="error")


if __name__ == "__main__":
    bot = TradingBot()
    bot.start_trading()


